Docs on hook call order #21120
Docs on hook call order #21120
Conversation
SkafteNicki
requested review from
Borda,
ethanwharris,
justusschock,
lantiga and
tchaton
as code owners
| trainer.fit() | ||
| │ | ||
| ├── setup(stage="fit") | ||
| │ └── [Callbacks only] |
There was a problem hiding this comment.
just a suggestion. perhaps adding this could also be helpful? as it also shows one of the known issue #19658
├── setup(stage="fit")
│ └── [Callback.setup]
│ └── [LightnintModule.configure_shared_model]
│ └── [LightnintModule.configure_model]
│ └── [Strategy.setup]
| │ └── [Strategy] | ||
| │ | ||
| ├── on_sanity_check_start() |
There was a problem hiding this comment.
perhaps we could also add when the checkpoint and state_dict are load?
on_fit_start()
if strategy.restore_checkpoint_after_setup:
[LightningModule.on_load_checkpoint]
[Strategy.load_model_state_dict()
[optimizer.load_state_dict]
[lr_scheduler.load_state_dict]
on_sanity_check_start()
| └── teardown(stage="fit") | ||
| └── [Callbacks only] |
There was a problem hiding this comment.
I believe that the lightning module also has a teardown callback here.
└── teardown(stage="fit")
└── [Callback]
└── [LightningModule]
| ├── on_fit_end() | ||
| │ ├── [Callbacks] | ||
| │ ├── [LightningModule] | ||
| │ └── [Strategy] | ||
| │ | ||
| └── teardown(stage="fit") |
There was a problem hiding this comment.
from my personal code inspect, I believe that the strategy.teardown should be called before on_fit_end stage, and also the strategy class does not have a on_fit_end callback.
strategy.teardown()
|--LightningModule.cpu()
on_fit_end()
|--Callback.on_fit_end()
|--LightningModule.on_fit_end()
Callback.teardown()
LightningModule.teardown()
| │ │ ├── on_before_zero_grad() | ||
| │ │ │ ├── [Callbacks] | ||
| │ │ │ └── [LightningModule] | ||
| │ │ │ | ||
| │ │ ├── [Forward Pass - training_step()] | ||
| │ │ │ └── [Strategy only] |
There was a problem hiding this comment.
I think the training step is called before the on_before_zero_grad is called?
strategy.training_step()
| -- LightningModule.training_step()
on_before_zero_grad()
|-- Callback.on_before_zero_grad()
|-- LightningModule.on_before_zero_grad()
|-- LightningModule.optimizer_zero_grad()
| │ │ ├── on_before_optimizer_step() | ||
| │ │ │ ├── [Callbacks] | ||
| │ │ │ └── [LightningModule] |
There was a problem hiding this comment.
perhaps also add configure_gradient_clipping here?
│ │ ├── on_before_optimizer_step()
│ │ │ ├── [Callbacks]
│ │ │ ├── [LightningModule]
│ │ │ └── [LightningModule.configure_gradient_clipping]
| │ ├── [Callbacks - Non-monitoring only] | ||
| │ ├── [LightningModule] | ||
| │ └── [Callbacks - Monitoring only] |
There was a problem hiding this comment.
Perhaps Non-monitoring could be changed to Checkpoint & EarlyStopping to be more specific.
|
@GdoongMathew, all good call. How about landing this as a baseline, and you can send just another PR with your suggestions for improvement? :) |
|
From my previous experience, since Lightning uses the closure function in the optimizer, it can make tracing the optimizer step a bit difficult. Maybe we could also highlight that here? but this also makes my personal note a bit messier as well.... |
|
@Borda sure!! will open another PR for those suggestions. thanks~ |
This comment was marked as off-topic.
This comment was marked as off-topic.
|
@GdoongMathew, it has landed, so now you can send your suggestions with improvments 🚀 |
* add hook order * add to index * Apply suggestions from code review * BoringModel * testoutput --------- Co-authored-by: Jirka Borovec <[email protected]> Co-authored-by: Jirka B <[email protected]> (cherry picked from commit 3d56296)
* add hook order * add to index * Apply suggestions from code review * BoringModel * testoutput --------- Co-authored-by: Jirka Borovec <[email protected]> Co-authored-by: Jirka B <[email protected]> (cherry picked from commit 3d56296)
3 participants
What does this PR do?
Fixes #21069
Before submitting
PR review
Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:
Reviewer checklist
📚 Documentation preview 📚: https://pytorch-lightning--21120.org.readthedocs.build/en/21120/